from Crypto.PublicKey import ECC
from Crypto.Signature import DSS
from Crypto.Hash import SHA256
from eth_account import Account
from eth_account.messages import encode_defunct
from web3 import Web3
from fl.logger import logger
from web3.exceptions import ContractLogicError
from fl.util.chain import BlockchainManager
import hashlib

class VRF(BlockchainManager):
    def __init__(self, config):
        super().__init__(config)
        self.vrf_account = Account.create() # 这边专门为vrf选举过程创建了临时的账户
        logger.info(f"VRF Account created: {self.vrf_account.address}")
    # 生成证明
    def generate_proof(self, seed):
        message_hash = Web3.keccak(seed)
        signed_message = self.vrf_account.signHash(message_hash)

        r = Web3.to_bytes(signed_message.r).rjust(32, b'\0')
        s = Web3.to_bytes(signed_message.s).rjust(32, b'\0')
        v = signed_message.v

        return r, s, v

    # 生成proof和随机数
    # def generate_random_number(self, seed):
    #     r, s, v, _ = self.generate_proof(seed)
    #     proof = (r, s, v)
    #     vrf_output = hashlib.sha256(r.to_bytes(32, byteorder='big') +
    #                                 s.to_bytes(32, byteorder='big') +
    #                                 v.to_bytes(1, byteorder='big')).hexdigest()
    #     return proof, vrf_output

    def register_on_chain(self):
        try:
            tx_hash = (self.contract.functions.register(self.vrf_account.address)
                       .transact({'from': self.account, 'nonce': self.w3.eth.get_transaction_count(self.account)}))
            self.w3.eth.wait_for_transaction_receipt(tx_hash)
        except ContractLogicError as e:
            logger.error(e.args[0])

    def submit_proof_on_chain(self, r, s, v):
        try:
            tx_hash = (self.contract.functions.submitProof(r, s, v).transact({'from': self.account,
                             'nonce': self.w3.eth.get_transaction_count(self.account)}))

            self.w3.eth.wait_for_transaction_receipt(tx_hash)
        except ContractLogicError as e:
            logger.error(e.args[0])

    def generate_seed_on_chain(self):
        try:
            tx_hash = (self.contract.functions.generateSeed()
                        .transact({'from': self.account,
                                   'nonce': self.w3.eth.get_transaction_count(self.account)}))

            self.w3.eth.wait_for_transaction_receipt(tx_hash)
        except ContractLogicError as e:
            logger.warning(e.args[0])

    def get_seed_from_chain(self):
        return self.contract.functions.getCurrentSeed().call()

    def select_nodes(self):
        try:
            tx_hash = (self.contract.functions.selectAggregator()
                       .transact({'from': self.account,
                                  'nonce': self.w3.eth.get_transaction_count(self.account)}))
            self.w3.eth.wait_for_transaction_receipt(tx_hash)
        except ContractLogicError as e:
            logger.error(e.args[0])

    def get_selected_nodes(self):
        return self.contract.functions.getSelectedAggregators().call()

    def close_registry_on_chain(self):
        try:
            tx_hash = (self.contract.functions.closeRegistry()
                        .transact({'from': self.account,
                                   'nonce': self.w3.eth.get_transaction_count(self.account)}))
            self.w3.eth.wait_for_transaction_receipt(tx_hash)
        except ContractLogicError as e:
            logger.error(e.args[0])


print(Account.create().address)